import os
import requests
from flask import Flask, request, jsonify
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import PyPDF2  # 用于读取 PDF 文件
from docx import Document  # 用于读取 Word 文件
import pandas as pd  # 用于读取 Excel 文件


def load_documents(folder_path):
    """加载文件夹中的所有文档，支持多种文件格式"""
    documents = []

    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)

        # 处理 .txt 文件
        if filename.endswith('.txt'):
            with open(file_path, 'r', encoding='utf-8') as file:
                documents.append(file.read())

        # 处理 .pdf 文件
        elif filename.endswith('.pdf'):
            with open(file_path, 'rb') as file:
                reader = PyPDF2.PdfReader(file)
                text = ""
                for page in reader.pages:
                    text += page.extract_text()
                documents.append(text)

        # 处理 .docx 文件
        elif filename.endswith('.docx'):
            doc = Document(file_path)
            text = ""
            for paragraph in doc.paragraphs:
                text += paragraph.text + "\n"
            documents.append(text)

        # 处理 .xlsx 文件
        elif filename.endswith('.xlsx'):
            df = pd.read_excel(file_path)
            text = df.to_string(index=False)  # 将表格内容转换为字符串
            documents.append(text)

        # 处理其他格式（可选）
        else:
            print(f"Unsupported file format: {filename}")

    return documents

def vectorize_documents(documents):
    """将文档向量化"""
    vectorizer = TfidfVectorizer()
    vectors = vectorizer.fit_transform(documents)
    return vectors, vectorizer

def retrieve_documents(query, vectorizer, vectors, documents, top_k=3):
    """根据用户查询检索相关文档"""
    query_vector = vectorizer.transform([query])
    similarities = cosine_similarity(query_vector, vectors).flatten()
    top_indices = similarities.argsort()[-top_k:][::-1]
    return [documents[i] for i in top_indices]

from openai import OpenAI

def call_deepseek_api(query, context):
    """通过 OpenAI SDK 调用 DeepSeek API"""    
    client = OpenAI(api_key="sk-0bb2e1837ff7453c873c259368dbdf8f", base_url="https://api.deepseek.com")

    try:
        # 发送请求
        messages = [{"role": "user", "content": f"{context}\n\n问题：{query}"}]
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=messages
        )

        # 打印调试信息
        print("Request Data:", messages)
        
        # 打印响应信息
        print("Response:", response)

        print(response.choices[0].message.content)
        
        # 返回生成的回答
        return response.choices[0].message.content
    except Exception as e:
        # 打印错误信息
        print("Error:", e)
        return None



app = Flask(__name__)

# 加载知识库
documents = load_documents(r'E:\\BijiApp\\web_app\\python_part\\Demo1\\TEMP')
vectors, vectorizer = vectorize_documents(documents)

@app.route('/query', methods=['POST'])
def query():
    # 获取用户查询
    user_query = request.json.get('query')
    
    # 检索相关文档
    relevant_docs = retrieve_documents(user_query, vectorizer, vectors, documents)
    context = "\n".join(relevant_docs)  # 将相关文档拼接为上下文
    
    # 调用 DeepSeek API
    deepseek_response = call_deepseek_api(user_query, context)
    
    # 返回结果
    return jsonify({
        'query': user_query,
        'context': context,
        'response': deepseek_response
    })

if __name__ == '__main__':
    app.run(port=5000)